/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.server;

import static org.sss.common.Utils.closeDataReceivers;
import static org.sss.common.Utils.newInstance;

import java.lang.reflect.Method;
import java.util.List;
import java.util.UUID;

import org.slf4j.Logger;
import org.sss.common.Bucket;
import org.sss.common.CanceledException;
import org.sss.common.CommonStatus;
import org.sss.common.DataReceiver;
import org.sss.common.ErrorDataReceiverException;
import org.sss.common.Inlet;
import org.sss.common.InputDrivenTask;
import org.sss.common.Reflection;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.StorageNode;
import org.sss.util.CallbackInterface;

public abstract class ProcTask<Proc, Context, Handled> extends TaskWithCounter implements InputDrivenTask<Handled>, RenewableTask {
  private static Logger logger = SssServer.getLogger(ProcTask.class);
  private final Reflection.ProcInfo procInfo;
  private final CallbackInterface<CommonStatus> doneCallback;
  private int numDoneSubTask;
  private Proc proc;
  private final ContextImpl contextImpl;
  private boolean configured = false;
  private boolean cleanup_done = false;
  private final int cacheSize;

  public ProcTask(TaskSet taskSet,
                  UUID jobID,
                  SssServerResource sssResource,
                  Configuration conf,
                  StorageNode storageNode,
                  Inlet inlet,
                  Class<? extends Proc> procClass,
                  List<org.sss.common.Outlet> outletList,
                  CallbackInterface<CommonStatus> doneCallback,
                  int cacheSize,
                  long queueMultiplicity,
                  long queueLimit) throws SssException {
    super(taskSet, jobID, sssResource, conf, storageNode, new Input(conf, inlet));
    if (inlet.GID.isBroadcast()) {
      throw new SssException(
          "Input group id has 'broadcast' type: " + inlet.GID);
    }
    this.doneCallback = doneCallback;
    this.contextImpl = new ContextImpl(sssResource, storageNode, taskSet, conf, this);
    this.procInfo = Reflection.createProcInfo(procClass);
    this.proc = newInstance(procClass);
    this.cacheSize = cacheSize;
    int arity = procInfo.getOutputTypes().size();
    assert arity == outletList.size(): String.format("mapper arity and outlet mismatch.(artiy = %d, size = %d)", arity, outletList.size());
    assert arity > 0:"output size is less than 1.";

    if (procInfo.getCleanupMethod() != null) {
      logger.info("task with cleanup method");
    }

    // TODO: run on another thread?
    for (org.sss.common.Outlet o: outletList) {
      getOutlets().create(o.GID,
          o.keyClassName, o.valueClassName, 
          o.combinerClassName, o.partitionerClassName,
          cacheSize, queueMultiplicity, queueLimit);
    }
  }

  @Override
  public void submit() throws SssException {
    this.numDoneSubTask = 0;
    getWKTP().execute(new Runnable() {
      @Override
      public void run() {
        try {
          if (!configured) {
            configure();
            configured = true;
          }
          submitReadTaskSlices();
        } catch (Throwable e){
          errorCallback(e);
        }
      }
    });
  }

  @Override
  public synchronized void noMoreTaskSlice(Throwable exp) {
    try {
      if (exp != null) {
        throw exp;
      }

      numDoneSubTask++;
      logger.debug("noMoreTaskSlice. {}", numDoneSubTask);
      if (numDoneSubTask >= getInputGID().getNSlot()) {
        getOutlets().closeAll();
      }
    } catch (ErrorDataReceiverException e) {
      logger.debug("Ignore ErrorDataReceiverException", e);
    } catch (CanceledException e) {
      getOutlets().setError(e);
    } catch (Throwable e) {
      logger.error("Catch an exception which is not handled.", e);
      getOutlets().setError(e);
    }
  }

  protected void doneCallback(CommonStatus status) {
    doneCallback.callback(status);
  }
  
  protected Reflection.ProcInfo getProcInfo() {
    return procInfo;
  }

  protected Proc getProc() {
    return proc;
  }

  /**
   * This call on completing to write tuples.
   */
  @Override
  public synchronized void ioWriteComplete(Throwable exp) {
    logger.info("ioWriteComplete.");
    final CommonStatus s = new CommonStatus();
    if (exp == null) {
      final Method cleanupMethod = procInfo.getCleanupMethod();
      if (cleanupMethod != null && !cleanup_done) {
        getWKTP().execute(new Runnable() {
          @Override
          public void run() {
            try {
              getOutlets().renew();
              List<DataReceiver<Bucket>> cache = getOutlets().getTLPutters(cacheSize);
              Object[] args = createArguments(createContext(), 0, cache);
              cleanupMethod.invoke(proc, args);
              closeDataReceivers(cache);
              getOutlets().closeAll();
            }
            catch (Throwable e) {
              s.exception = e;
              doneCallback(s);
            }
          }
        });
        cleanup_done = true;
      }
      else  {
        setValueToStatus(s, exp);
        doneCallback(s);
      }
    }
    else {
      cancel();
      s.exception = exp;
      doneCallback(s);
    }
  }

  protected Context createContext() {
    return createContextImpl(contextImpl);
  }

  protected synchronized void errorCallback(Throwable e) {
    CommonStatus s = new CommonStatus();
    s.exception = e;
    doneCallback(s);
  }

  protected abstract void configure() throws Exception;
  protected abstract void submitReadTaskSlices() throws SssException;
  protected abstract Context createContextImpl(ContextImpl impl);
}
